from random import randint

import json
import pandas as pd
import os
import pytorch_lightning as pl
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize, RandomCrop

try:
    from torchvision.transforms import InterpolationMode

    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC

def _convert_image_to_rgb(image):
    return image.convert("RGB")

def _transform(img_px, crop_px):
    return Compose(
        [
            _convert_image_to_rgb,
            Resize(img_px, interpolation=BICUBIC),
            RandomCrop(crop_px),
            ToTensor(),
        ]
    )


class ImageDataset(Dataset):
    def __init__(
        self,
        df_path,
        image_size=512,
        crop_size=256,
        infinity=False,
    ):
        with open(df_path, 'r', encoding='utf-8') as f:
            img_dir = os.path.dirname(df_path)
            self.imgs = [os.path.join(img_dir, x) for x in list(json.loads(f.read()).keys())]
        self.image_size = image_size
        self.infinity = infinity
        self.transform =_transform(image_size, crop_size)

    def __len__(self):
        if self.infinity:
            return 99999999
        else:
            return len(self.imgs)

    def __getitem__(self, item):
        if self.infinity:
            ind = randint(0, len(self.imgs)-1)
        else:
            ind = item
        image = Image.open(self.imgs[ind])
        image = self.transform(image)
        image = image*2.-1.

        return image

def create_loader(batch_size, num_workers, shuffle=False, **dataset_params):
    dataset = ImageDataset(**dataset_params)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        shuffle=shuffle,
        pin_memory=True,
    )

class LightningDataModule(pl.LightningDataModule):
    """PyTorch Lightning data class"""

    def __init__(self, train_config):
        super().__init__()
        self.train_config = train_config

    def train_dataloader(self):
        return create_loader(**self.train_config)
